#!/bin/bash

if [[ -f /etc/lab/lab-functions ]]; then
    . /etc/lab/lab-functions
fi

#boot /dev/sda1 100M
#root1 /dev/sda2 rootfs size * 1.2, around 600M
#root2 /dev/sda3 rootfs size * 1.2
#swap /dev/sda5 10% of disk max 4G
#var /dev/sda6 20% of disk 
#app1 /dev/sda7 10% max 1G
#app2 /dev/sda8 10% max 1G
#data /dev/sda9 remaining disk

#create_parts as above
create_parts()
{
    hd=$1

    pn=$(fdisk -l $hd | wc -l)
    if [[ $pn -gt 15 ]]; then
        echo "Partitions already created"
        return 1
    fi

    disk_size=$(sfdisk -s $hd)
    disk_size=$(( $disk_size/1024 ))

    boot_size=100 #MB

    #required root partition size
    factor=5
    root_size=$(du -mx / | tail -1 | awk '{printf("%d", $1)}')
    root_size=$(printf %.0f $(echo "$root_size*$factor" | /usr/bin/bc))

    if [[ $(( $root_size*2 )) -gt $disk_size ]]; then
        echo "Not enough disk space ($disk_size MB) required (2 * $root_size MB), quitting"
        exit 1
    fi

    #swap size
    swap_size=$(printf %.0f $(echo "$disk_size*0.1" | /usr/bin/bc))
    if [[ $swap_size -gt 4000 ]]; then
        swap_size=4000
    fi

    # var size
    var_size=$(printf %.0f $(echo "$disk_size*0.2" | /usr/bin/bc))

    # app size
    app_size=$(printf %.0f $(echo "$disk_size*0.1" | /usr/bin/bc))

    echo -n "Formatting disk $hd ($disk_size MB) ... "
    dd if=/dev/zero of=$hd bs=100k count=1024 >/dev/null 2>&1
    echo "Done"

    create_partition ${hd} 1 primary $boot_size boot 
    create_partition ${hd} 2 primary $root_size root1 
    create_partition ${hd} 3 primary $root_size root2 
    create_partition ${hd} 4 extended 0 ""
    create_partition ${hd} 5 logical $swap_size swap
    create_partition ${hd} 6 logical $var_size var 
    create_partition ${hd} 7 logical 0 data 
}

format_parts()
{
    echo -n "Formatting partitions ... "
    hd=$1
    dd if=/dev/zero of=${hd}1 bs=100k count=512 >/dev/null 2>&1
    dd if=/dev/zero of=${hd}2 bs=100k count=512 >/dev/null 2>&1
    dd if=/dev/zero of=${hd}3 bs=100k count=512 >/dev/null 2>&1
    dd if=/dev/zero of=${hd}6 bs=100k count=512 >/dev/null 2>&1
    echo "Done"

    mk_fs ${hd}1 ext3 boot
    mk_fs ${hd}2 ext3 root1
    mk_fs ${hd}3 ext3 root2
    mk_fs ${hd}6 ext3 var
}

mount_parts()
{
    echo -n "Mounting partitions ... "
    hd=$1
    mkdir -p /mnt/{boot,root1,root2,var,app1,app2}
    mount ${hd}1 /mnt/boot
    mount ${hd}2 /mnt/root1
    mount ${hd}3 /mnt/root2
    mount ${hd}6 /mnt/var
    echo "Done"
}

umount_parts()
{
    echo -n "Un-mounting partitions ... "
    hd=$1
    umount ${hd}1 /mnt/boot 2>/dev/null
    umount ${hd}2 /mnt/root1 2>/dev/null
    umount ${hd}3 /mnt/root2 2>/dev/null
    umount ${hd}6 /mnt/var 2>/dev/null
    echo "Done"
}

#create the fstab
create_fstab()
{
    hd=$1
    cat << EOF > /tmp/fstab
${hd}2                  /                       ext3    defaults        0 1
${hd}1                  /boot                   ext3    defaults        0 0
${hd}6                  /var                    ext3    defaults        0 0
tmpfs                   /dev/shm                tmpfs   defaults        0 0
devpts                  /dev/pts                devpts  gid=5,mode=620  0 0
sysfs                   /sys                    sysfs   defaults        0 0
proc                    /proc                   proc    defaults        0 0
${hd}5                  swap                    swap    defaults        0 0
EOF
}

copy_rootfs()
{
    #cp -ax is not keeping it to same filesystem, hence using tar
    cd $1
    tar -cf - --one-file-system --exclude /var / | tar -xf - 2>/dev/null
    #copy essential dev
    cp -a /dev/{MAKEDEV,null,console,zero,urandom,random} $1/dev/
    mkdir -p $1/dev/{shm,pts}
    mkdir -p $1/var
    cp -a /tmp/fstab $1/etc/
}

#Do the grub thing
create_grub()
{
    hd=$1
    kernel=$(basename /mnt/boot/vmlinuz*)
    cp -a /usr/share/grub/*/* /mnt/boot/grub/

    cat << EOF > /mnt/boot/grub/grub.conf
default=0
timeout=5
title LAB Active
        root (hd0,0)
        kernel /$kernel ro root=${hd}2
        initrd /initrd.gz

title LAB Backup
        root (hd0,0)
        kernel /$kernel ro root=${hd}3
        initrd /initrd2.gz
EOF

echo " Installing grub on primary partition..."
grub --no-floppy --batch <<EOF >/dev/null 2>&1
root (hd0,0)
setup (hd0)
quit
EOF

}

# ----------------------- MAIN ------------------------
disk=$1
umount_parts $disk
create_parts $disk
format_parts $disk
if [[ $? -ne 0 ]]; then
    echo "Couldn't format partitions on $disk"
    exit
fi
mount_parts $disk
create_fstab $disk

cp -a /boot/* /mnt/boot/
copy_rootfs /mnt/root1
#clear the buffer cache to free up some memory
sysctl -w vm.drop_caches=3
copy_rootfs /mnt/root2
cp -a /var/* /mnt/var

echo " Creating initrd image ..."
/sbin/mkinitrd --fstab=/mnt/root1/etc/fstab /mnt/boot/initrd.gz `uname -r` >/dev/null 2>&1
/sbin/mkinitrd --fstab=/mnt/root2/etc/fstab /mnt/boot/initrd2.gz `uname -r` >/dev/null 2>&1

echo " Setting up grub ..."
create_grub $disk
